import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from statsmodels.nonparametric.smoothers_lowess import lowess
from matplotlib.dates import DateFormatter
from datetime import timedelta


def loess_future_prediction_save(file_path="output.xlsx", frac=0.07, future_days=90, window_size=30,
                                  confidence_interval_alpha=0.05, output_file="result.xlsx"):
    """
    Reads data, performs LOESS smoothing, predicts future data, and calculates confidence intervals
    using a sliding window on residuals. Saves all data, predictions, and confidence intervals to an Excel file.
    All data is used as training data.

    Args:
        file_path (str, optional): Path to the Excel file. Defaults to "output.xlsx".
        frac (float, optional): LOESS smoothing parameter. Defaults to 0.2.
        future_days (int, optional): Number of days to predict into the future. Defaults to 90.
        window_size (int, optional): Size of the sliding window for confidence intervals. Defaults to 30.
        confidence_interval_alpha (float, optional): Confidence level. Defaults to 0.05.
        output_file (str, optional): Path to the output Excel file. Defaults to "result.xlsx".

    Returns:
        None: Generates and displays a visualization plot, and saves data to an Excel file.
    """

    # 1. Read data
    df = pd.read_excel(file_path)
    df['Date'] = pd.to_datetime(df['Date'])
    df = df.set_index('Date')
    y = df.iloc[:, 2].values
    x = df.index.values

    # 2. Use all data as training data
    x_train = x.copy()
    y_train = y.copy()

    # 3. Iterative Prediction
    future_dates = []
    future_values = []
    last_date = pd.to_datetime(x[-1])

    # Convert x_train to datetime objects
    current_x_train = [pd.to_datetime(date).to_pydatetime() for date in x_train]
    current_y_train = y_train.copy()  # Create a copy of y_train for iterative updates

    for i in range(1, future_days + 1):
        # LOESS smoothing - Use updated training data
        # Convert timestamps to numerical values
        current_x_train_num = [date.timestamp() for date in current_x_train]
        loess_smoothed = lowess(current_y_train, current_x_train_num, frac=frac)
        smoothed_x = loess_smoothed[:, 0]
        smoothed_y = loess_smoothed[:, 1]

        # Predict one future day
        next_date = last_date + timedelta(days=i)
        future_dates.append(next_date)

        # Interpolation prediction
        unique_x = np.unique(smoothed_x)
        unique_y = smoothed_y[np.isin(smoothed_x, unique_x)]
        future_x_num = [next_date.timestamp()]
        next_value = np.interp(future_x_num, unique_x, unique_y)[0]

        future_values.append(next_value)

        # Update training set
        current_x_train = np.append(current_x_train, pd.to_datetime(next_date).to_pydatetime())
        current_y_train = np.append(current_y_train, next_value)

    future_dates = np.array(future_dates)
    future_values = np.array(future_values)

    # 4. LOESS Smoothing on training data
    smoothed_x_train = lowess(y_train, x_train, frac=frac)[:, 0]
    smoothed_dates_train = [pd.to_datetime(date) for date in smoothed_x_train]
    smoothed_y_train = lowess(y_train, x_train, frac=frac)[:, 1]


    # 5. Calculate residuals and confidence intervals
    residuals = y_train - smoothed_y_train

    def calculate_confidence_interval(data, window_size, alpha):
        lower_bounds = []
        upper_bounds = []
        z = abs(np.percentile(np.random.randn(1000), alpha / 2 * 100))

        for i in range(len(data)):
            start = max(0, i - window_size + 1)
            window = data[start:i + 1]
            std = np.std(window)
            lower = -z * std
            upper = z * std
            lower_bounds.append(lower)
            upper_bounds.append(upper)
        return lower_bounds, upper_bounds

    train_lower_residual, train_upper_residual = calculate_confidence_interval(residuals, window_size,
                                                                         confidence_interval_alpha)
    train_lower = smoothed_y_train + train_lower_residual
    train_upper = smoothed_y_train + train_upper_residual


    # Calculate confidence intervals for future values (using the last residual std)
    forecast_ci_lower = []
    forecast_ci_upper = []

    last_residual_std = np.std(residuals[-window_size:]) if len(residuals) >= window_size else np.std(residuals)
    z = abs(np.percentile(np.random.randn(1000), confidence_interval_alpha / 2 * 100))
    forecast_lower_residual = -z * last_residual_std
    forecast_upper_residual = z * last_residual_std
    for value in future_values:
        forecast_ci_lower.append(value + forecast_lower_residual)
        forecast_ci_upper.append(value + forecast_upper_residual)


    # 6. Prepare data for saving
    train_df = pd.DataFrame({'Date': x_train, 'Original': y_train})
    train_df['LOESS Smoothed'] = smoothed_y_train
    train_df['Training CI Lower'] = train_lower
    train_df['Training CI Upper'] = train_upper

    forecast_df = pd.DataFrame({'Date': future_dates, 'Prediction': future_values})
    forecast_df['Prediction CI Lower'] = forecast_ci_lower
    forecast_df['Prediction CI Upper'] = forecast_ci_upper


    result_df = pd.concat([train_df, forecast_df], ignore_index=True)
    result_df['Date'] = pd.to_datetime(result_df['Date'])

    # 7. Save to Excel
    result_df.to_excel(output_file, index=False)

    # 8. Visualize
    fig, ax = plt.subplots(figsize=(14, 7))
    ax.plot(train_df['Date'], train_df['Original'], label="Training Data (Original)", marker='.', linestyle='-',
            alpha=0.5)
    ax.plot(train_df['Date'], train_df['LOESS Smoothed'], label="Training Data (LOESS Smoothed)", color='red')
    ax.plot(forecast_df['Date'], forecast_df['Prediction'], label=f"Future {future_days} Days (Predicted)", color='blue',
            linestyle='--')

    # Plot confidence intervals for training data
    ax.fill_between(train_df['Date'], train_df['Training CI Lower'], train_df['Training CI Upper'], color='red',
                    alpha=0.2,
                    label='Training Confidence Interval')
    # Plot confidence intervals for predicted data
    ax.fill_between(forecast_df['Date'], forecast_df['Prediction CI Lower'], forecast_df['Prediction CI Upper'],
                    color='blue', alpha=0.2,
                    label='Prediction Confidence Interval')

    # Set date format
    date_format = DateFormatter("%Y-%m-%d")
    ax.xaxis.set_major_formatter(date_format)
    fig.autofmt_xdate()

    ax.set_xlabel("Date")
    ax.set_ylabel("Data Value")
    ax.set_title("LOESS Smoothing, Future Prediction, and Confidence Intervals (Residual STD)")
    ax.legend()
    plt.grid(True)
    plt.show()


if __name__ == "__main__":
    loess_future_prediction_save()